package com.aiwiown.snackmq.storage.transaction;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.serialization.JsonSerializer;
import com.aiwiown.snackmq.common.serialization.SerializerFactory;
import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 事务日志，用于存储半消息及其状态。
 * <p>
 * 【已重构】: 此版本统一了内存缓存，解决了数据类型不匹配的问题，并简化了内部逻辑。
 * </p>
 */
@Slf4j
public class TransactionLog implements Closeable {

    private final File logFile;
    private final RandomAccessFile raf;
    private final ReentrantLock lock = new ReentrantLock();
    private final JsonSerializer serializer = SerializerFactory.getJsonSerializer();

    /**
     * 【已重构】内存索引，存储所有处于 PREPARE 状态的事务记录。
     * Key: transactionId
     * Value: 完整的 TransactionRecord 对象
     */
    private final Map<String, TransactionRecord> pendingTransactions = new ConcurrentHashMap<>();

    public TransactionLog(String path) throws IOException {
        File dir = new File(path, "transaction");
        if (!dir.exists() && !dir.mkdirs()) {
            throw new IOException("Failed to create transaction log directory: " + dir.getAbsolutePath());
        }
        this.logFile = new File(dir, "transaction.log");
        this.raf = new RandomAccessFile(logFile, "rw");
        recover();
    }

    /**
     * 从日志文件中恢复所有待处理的事务。
     * 启动时调用，以重建内存索引。
     */
    private void recover() {
        lock.lock();
        try {
            raf.seek(0);
            long currentPosition = 0;
            while (currentPosition < raf.length()) {
                // 1. 读取记录长度
                if (raf.length() - currentPosition < 4) {
                    log.warn("Found partial record header at position {}. Truncating transaction log.", currentPosition);
                    break;
                }
                int recordSize = raf.readInt();
                if (recordSize <= 0) {
                    log.warn("Found invalid record size ({}) at position {}. Truncating transaction log.", recordSize, currentPosition);
                    break;
                }

                // 2. 读取记录体
                if (raf.length() - currentPosition - 4 < recordSize) {
                    log.warn("Found partial record body at position {}. Truncating transaction log.", currentPosition);
                    break;
                }
                byte[] recordBytes = new byte[recordSize];
                raf.readFully(recordBytes);

                // 3. 反序列化并处理记录
                TransactionRecord record = serializer.deserialize(recordBytes, TransactionRecord.class);
                if (record != null) {
                    // 【修改】调用新的 processRecord，不再需要 position 参数
                    processRecord(record);
                }

                currentPosition = raf.getFilePointer();
            }
            // 截断损坏的文件部分
            raf.setLength(currentPosition);
            // 【修改】日志记录使用新的缓存
            log.info("TransactionLog recovery complete. Found {} pending transactions.", pendingTransactions.size());
        } catch (IOException e) {
            log.error("Failed to recover transaction log. State might be inconsistent.", e);
        } finally {
            lock.unlock();
        }
    }

    /**
     * 将一条事务记录追加到日志末尾。
     *
     * @param record 要追加的记录
     * @throws IOException 如果写入失败
     */
    public void append(TransactionRecord record) throws IOException {
        lock.lock();
        try {
            long position = raf.length();
            byte[] recordBytes = serializer.serialize(record);
            raf.seek(position);
            raf.writeInt(recordBytes.length);
            raf.write(recordBytes);

            // 【修改】更新内存索引
            processRecord(record);
        } finally {
            lock.unlock();
        }
    }

    /**
     * 【已重构】根据记录类型更新内存索引。
     *
     * @param record 事务记录
     */
    private void processRecord(TransactionRecord record) {
        String txId = record.getTransactionId();
        switch (record.getRecordType()) {
            case PREPARE:
                pendingTransactions.put(txId, record);
                break;
            case COMMITTED:
            case ROLLED_BACK:
                pendingTransactions.remove(txId);
                break;
            default:
                log.warn("Unknown transaction record type: {}", record.getRecordType());
        }
    }

    /**
     * 获取一个待处理的半消息。
     *
     * @param transactionId 事务ID
     * @return 半消息对象，如果不存在则返回 null
     */
    public Message getPendingMessage(String transactionId) {
        // 【修改】从缓存的 TransactionRecord 中获取消息
        TransactionRecord record = pendingTransactions.get(transactionId);
        return (record != null) ? record.getMessage() : null;
    }

    /**
     * 【已重构】获取所有待处理的事务记录。
     *
     * @return 一个包含待处理事务记录的 Map
     */
    public Map<String, TransactionRecord> getPendingTransactions() {
        // 【修改】返回一个防御性副本，防止外部修改
        return new HashMap<>(pendingTransactions);
    }

    @Override
    public void close() throws IOException {
        lock.lock();
        try {
            if (raf != null) {
                raf.close();
            }
        } finally {
            lock.unlock();
        }
    }
}